#include <hal/pci.h>
#include <hal/apic.h>
#include <mm/valloc.h>

#define LOG_TAG "PCI"
#include <log.h>

lnode_t pci_devices;

void pci_probe_msi(pci_dev_t* dev) {
    u32 status = pci_read_cspace(dev->cspace_base, 0x4) >> 16;
    if (!(status & 0x10)) {
        dev->msi_loc = 0;
        return;
    }
    u32 cap_ptr = pci_read_cspace(dev->cspace_base, 0x34);
    u32 cap_id;
    while (cap_ptr) {
        cap_id = pci_read_cspace(dev->cspace_base, cap_ptr);
        if ((cap_id & 0xFF) == 0x5) {
            dev->msi_loc = cap_ptr;
            return;
        }
        cap_ptr = (cap_id >> 8) & 0xFF;
    }
}

void pci_probe_device(int bus, int dev, int func) {
    pci_cfg_addr_reg_t addr = {
        .enable = 1,
        .bus_num = bus,
        .dev_num = dev,
        .func_num = func
    };
    u32 devId_venId = pci_read_cspace(addr, 0);
    if ((devId_venId & 0xFFFF) == 0xFFFF)
        return;

    u32 header_type = pci_read_cspace(addr, 0xC);
    header_type = (header_type >> 16) & 0xFF;
    if ((header_type & PCI_TYPE_MULTIFUNCTION) && func == 0) {
        header_type = header_type & ~PCI_TYPE_MULTIFUNCTION;
        // multi-function device
        for (int i = 1; i < 7; i++) {
            pci_probe_device(bus, dev, i);
        }
    }

    if (header_type != PCI_TYPE_DEVICE)
        return;

    u32 class = pci_read_cspace(addr, 0x8);
    u32 intr = pci_read_cspace(addr, 0x3C);

    pci_dev_t* device = valloc(sizeof(pci_dev_t));
    *device = (pci_dev_t){
        .intr_info = intr,
        .class_info = class,
        .cspace_base = addr,
        .device_info = devId_venId
    };

    pci_probe_msi(device);

    list_append(&pci_devices, &device->list);
}

void print_pci_devices() {
    pci_dev_t* pos, * n;
    list_for_each(pos, n, &pci_devices, list) {
        printk("[PCI] (B%xh:D%xh:F%xh) Dev %x:%x, Class 0x%x\n",
            pos->cspace_base.bus_num,
            pos->cspace_base.dev_num & 0x1F,
            pos->cspace_base.func_num,
            pos->device_info & 0xFFFF,
            pos->device_info >> 16,
            pos->class_info >> 8
        );
    }
}

void pci_init(void) {
    list_head_init(&pci_devices);
    for (int bus = 0; bus < 2; bus++) {
        for (int dev = 0; dev < 32; dev++) {
            pci_probe_device(bus, dev, 0);
        }
    }
    // print_pci_devices();
}

pci_dev_t* pci_get_by_class(u32 class) {
    pci_dev_t* pos, * n;
    list_for_each(pos, n, &pci_devices, list) {
        if ((pos->class_info >> 8) == class) {
            return pos;
        }
    }
    LOGW("No device found with class 0x%X\n", class);
    return NULL;
}

pci_dev_t* pci_get_by_id(u16 vendorId, u16 deviceId) {
    pci_dev_t* pos, * n;
    list_for_each(pos, n, &pci_devices, list) {
        if ((pos->device_info & 0xFFFF) == vendorId &&
            (pos->device_info >> 16) == deviceId) {
            return pos;
        }
    }
    LOGW("No device found with id: %X:%X\n", vendorId, deviceId);
    return NULL;
}

size_t pci_bar_sizing(pci_dev_t* dev, u32* bar_out, u32 bar_num) {
    u32 bar = pci_read_cspace(dev->cspace_base, PCI_REG_BAR(bar_num));
    if (!bar) {
        *bar_out = 0;
        return 0;
    }

    pci_write_cspace(dev->cspace_base, PCI_REG_BAR(bar_num), 0xFFFFFFFF);
    u32 size = pci_read_cspace(dev->cspace_base, PCI_REG_BAR(bar_num)) & ~0x1;
    if (PCI_BAR_MMIO(bar)) {
        size = PCI_BAR_ADDR_MM(size);
    }
    *bar_out = bar;
    pci_write_cspace(dev->cspace_base, PCI_REG_BAR(bar_num), bar);
    return ~size + 1;
}

void pci_setup_msi(pci_dev_t* dev, int vector) {
    u32 msi_addr = APIC_BASE_PA;
    u32 msi_data = vector;

    pci_write_cspace(dev->cspace_base, PCI_MSI_ADDR(dev->msi_loc), msi_addr);
    u32 reg = pci_read_cspace(dev->cspace_base, dev->msi_loc);
    u32 msi_ctl = reg >> 16;

    int offset = !!(msi_ctl & MSI_CAP_64BIT) * 4;
    pci_write_cspace(dev->cspace_base, PCI_MSI_DATA(dev->msi_loc, offset), msi_data & 0xFFFF);

    if (msi_ctl & MSI_CAP_MASK) {
        pci_write_cspace(dev->cspace_base, PCI_MSI_MASK(dev->msi_loc, offset), 0);
    }

    reg = (reg & 0xFF8FFFFF) | 0x10000;
    pci_write_cspace(dev->cspace_base, dev->msi_loc, reg);
}